// Copyright (c) 2019 Presto Labs Pte. Ltd.
// Author: jaewon

#include "coin1/fastfeed/feed_cache.h"

#include <string>
#include <utility>

#include <fmt/printf.h>
#include <glog/logging.h>
#include <boost/algorithm/string.hpp>

#include "coin1/base/exception_util.h"
#include "coin2/feed/fastfeed/reader.h"

namespace fastfeed {
namespace detail {

ReaderContext::ReaderContext(
    const std::vector<std::experimental::filesystem::path>& filepaths,
    std::function<void(const proto::Feed&)> callback)
    : filepaths_(filepaths),
      callback_(std::move(callback)),
      feed_(),
      buf_(new char[max_record_size]) {
  std::vector<std::unique_ptr<orio::io::Reader<proto::Feed>>> readers;

  for (auto& path : filepaths_) {
    std::unique_ptr<orio::io::Reader<orio::record::SimpleRecord>> reader(
        orio::record::SimpleReader::FromFile(path, max_record_size));
    std::unique_ptr<orio::io::Reader<proto::Feed>> fastfeed_reader(
        new FastFeedProtoReader(std::move(reader), max_record_size));
    readers.push_back(std::move(fastfeed_reader));
  }

  reader_.reset(
      new MergedFastFeedReader(std::move(readers)));
}

bool ReaderContext::ReadNextAndParse(proto::Feed* feed) {
  return reader_->Read(feed);
}

bool ReaderContext::ReadNext() { return ReadNextAndParse(&feed_); }

FrontRunReaderContext::FrontRunReaderContext(
    const std::vector<std::experimental::filesystem::path>& filepaths,
    std::function<void(const proto::Feed&)> callback)
    : ReaderContext(filepaths, callback), feed_(&feeds_[0]), feed_next_(&feeds_[1]) {}

bool FrontRunReaderContext::ReadNext() {
  if (!next_ready_) {
    next_ready_ = true;
    ReadNext();
  }

  // Swaps next and current.
  std::swap(feed_next_, feed_);
  result_ = result_next_;

  // Reads next.
  result_next_ = ReadNextAndParse(feed_next_);

  // Fills feed_continue.
  if (result_ && result_next_ &&
      (feed_->feed_continue() == proto::FEED_CONTINUE_INVALID ||
       feed_->feed_continue() == proto::FEED_CONTINUE_UNKNOWN)) {
    if (feed_->timestamp() == feed_next_->timestamp()) {
      feed_->set_feed_continue(proto::FEED_CONTINUE_YES);
    } else {
      feed_->set_feed_continue(proto::FEED_CONTINUE_NO);
    }
  }

  return result_;
}

}  // namespace detail

std::string FeedCacheKey::ToString() const {
  return fmt::format("{}/{}/W{}", machine.c_str(), sub_req.c_str(), worker_id.c_str());
}

FeedCacheReader::FeedCacheReader(
    const std::string& date_str,
    const std::vector<filesystem::path>& feed_cache_dirs)
    : date_str_(date_str), feed_cache_dirs_(feed_cache_dirs) {}

void FeedCacheReader::AddCacheKey(
    const FeedCacheKey& feed_cache_key,
    const std::vector<std::string>& symbols,
    std::function<void(const proto::Feed&)> callback) {
  for (auto&& symbol : symbols) {
    auto filepaths = ResolveCacheFilePath(feed_cache_key, symbol);
    if (filepaths.empty()) {
      LOG(WARNING) << "Feed cache for " << feed_cache_key.ToString() << ", "
                   << date_str_ << ", " << symbol
                   << " does not exist. Skipping";
      continue;
    }

    std::shared_ptr<detail::ReaderContext> reader(
        new detail::FrontRunReaderContext(filepaths, callback));
    PushOrDeleteReader(reader);
  }
}

bool FeedCacheReader::Read() {
  auto reader = PopReader();
  if (!reader) {
    return false;
  }
  reader->InvokeCallback();
  PushOrDeleteReader(reader);
  return true;
}

std::optional<int64_t> FeedCacheReader::GetNextTimestamp() const {
  if (pqueue_.empty()) {
    return std::nullopt;
  }
  auto reader = pqueue_.top();
  return reader->timestamp();
}

std::vector<std::experimental::filesystem::path>
FeedCacheReader::ResolveCacheFilePath(const FeedCacheKey& cache_key,
                                      const std::string& symbol) {
  std::set<std::string> extension = {"lz4", "gz"};
  std::string w = fmt::format("W{}", cache_key.worker_id.c_str());
  std::vector<std::experimental::filesystem::path> out;
  for (const filesystem::path& dir : feed_cache_dirs_) {
    filesystem::path filepath = dir;
    filepath.append(cache_key.machine)
        .append(date_str_)
        .append(cache_key.sub_req);
    if (!std::experimental::filesystem::exists(filepath)) continue;
    for (auto& path : std::experimental::filesystem::directory_iterator(filepath)) {
      auto filename = path.path().filename().string();
      if (filename.rfind(symbol) == 0) {
        std::vector<std::string> splitted;
        boost::split(splitted, filename, boost::is_any_of("."));
        CHECK_GE(splitted.size(), 2);
        std::string worker;
        if (extension.count(splitted.back()) > 0) {
          worker = splitted[splitted.size() - 2];
        } else {
          worker = splitted.back();
        }

        if (worker == w) {
          out.push_back(path);
        }
      }
    }

    if (!out.empty()) break;  // found in 1 dir
  }
  return out;
}

void FeedCacheReader::PushOrDeleteReader(std::shared_ptr<detail::ReaderContext> reader) {
  auto res = reader->ReadNext();
  if (res) {
    pqueue_.push(reader);
  }
}

std::shared_ptr<detail::ReaderContext> FeedCacheReader::PopReader() {
  if (pqueue_.empty()) {
    return std::shared_ptr<detail::ReaderContext>();
  }
  auto reader = pqueue_.top();
  pqueue_.pop();
  return reader;
}

FeedCacheRunner::FeedCacheRunner(
    std::function<void(const FeedCacheKey&, FeedBuilderBase*)> feed_builder_reset_callback,
    const std::string& date_str,
    const std::vector<filesystem::path>& feed_cache_dirs)
    : feed_builder_reset_callback_(feed_builder_reset_callback),
      date_str_(date_str),
      feed_cache_dirs_(feed_cache_dirs) {}

void FeedCacheRunner::AddCacheKey(
    const FeedCacheKey& feed_cache_key,
    const std::vector<std::string>& symbols) {
  if (feed_cache_key.worker_id.empty()) {
    AddCacheKeyWithFeedArb(feed_cache_key.machine, feed_cache_key.sub_req, symbols);
    return;
  }

  DCHECK_THROW(!feed_cache_key.worker_id.empty(), std::logic_error);
  FeedCacheKey key = feed_cache_key;
  key.worker_id.clear();

  auto& holder = req_[key];
  holder.symbols.insert(symbols.begin(), symbols.end());
  holder.worker_ids.insert(feed_cache_key.worker_id);
}

void FeedCacheRunner::AddCacheKey(const FeedCacheSubReq& request) {
  AddCacheKey(request.feed_cache_key, request.symbols);
}

void FeedCacheRunner::AddCacheKeyWithFeedArb(
    const std::string& machine,
    const std::string& sub_req,
    const std::vector<std::string>& symbols) {
  AddCacheKey(FeedCacheKey {machine, "1", sub_req}, symbols);
  AddCacheKey(FeedCacheKey {machine, "2", sub_req}, symbols);
}

void FeedCacheRunner::Run() { Run(std::nullopt, std::nullopt); }

void FeedCacheRunner::Run(
    std::optional<int64_t> from_timestamp,
    std::optional<int64_t> to_timestamp) {
  FeedCacheReader reader(date_str_, feed_cache_dirs_);

  for (auto& key_value : req_) {
    auto& holder = key_value.second;
    auto feed_builder =
        new DoubleShiftFeedBuilder(holder.worker_ids.size(), key_value.first.sub_req);
    feed_builder->set_from_timestamp(from_timestamp);
    holder.feed_builder.reset(feed_builder);

    std::vector<std::string> symbols(holder.symbols.begin(), holder.symbols.end());

    int worker_idx = 0;
    for (auto&& worker_id : holder.worker_ids) {
      FeedCacheKey cache_key = key_value.first;
      cache_key.worker_id = worker_id;
      reader.AddCacheKey(cache_key, symbols, feed_builder->GetUpdateFunction(worker_idx));
      ++worker_idx;
    }
    feed_builder_reset_callback_(key_value.first, feed_builder);
  }

  if (to_timestamp.has_value()) {
    while (true) {
      auto ts = reader.GetNextTimestamp();
      if (ts.has_value() && ts.value() >= to_timestamp.value()) {
        break;
      }
      if (!reader.Read()) {
        break;
      }
    }
  } else {
    while (reader.Read()) {
    }
  }
}

}  // namespace fastfeed
